#include <iostream>
#include <string>
#include <boost/asio.hpp>
#include <boost/beast.hpp>
#include <nlohmann/json.hpp>
#include <chrono>
#include <thread>

namespace beast = boost::beast;
namespace http = beast::http;
namespace websocket = beast::websocket;
namespace net = boost::asio;
using tcp = boost::asio::ip::tcp;
using json = nlohmann::json;

// 消息发送函数
bool send_message(websocket::stream<tcp::socket>& ws, const json& message) {
    try {
        std::string message_str = message.dump();
        ws.write(net::buffer(message_str));
        std::cout << "Sent message: " << message_str << std::endl;
        return true;
    } catch (const std::exception& e) {
        std::cerr << "Send error: " << e.what() << std::endl;
        return false;
    }
}

// 消息接收函数，带超时处理
bool receive_message(websocket::stream<tcp::socket>& ws, std::string& response, int timeout_seconds = 5) {
    try {
        // 设置读取操作的超时
        // ws.next_layer().set_option(boost::asio::socket_base::receive_timeout(std::chrono::seconds(timeout_seconds)));
        
        beast::flat_buffer buffer;
        ws.read(buffer);
        
        response = beast::buffers_to_string(buffer.data());
        std::cout << "Received response: " << response << std::endl;
        return true;
    } catch (const std::exception& e) {
        std::cerr << "Receive error: " << e.what() << std::endl;
        return false;
    }
}

int main(int argc, char* argv[]) {
    try {
        // 检查命令行参数
        std::string host = "127.0.0.1";
        unsigned short port = 8080;

        // 如果有命令行参数，则使用参数指定的host和port
        if (argc >= 2) {
            host = argv[1];
        }
        if (argc >= 3) {
            port = static_cast<unsigned short>(std::stoi(argv[2]));
        }

        std::cout << "Connecting to " << host << ":" << port << std::endl;

        // 创建I/O上下文
        net::io_context ioc;

        // 创建TCP解析器
        tcp::resolver resolver{ioc};

        // 创建WebSocket流
        websocket::stream<tcp::socket> ws{ioc};

        // 查找域名
        auto const results = resolver.resolve(host, std::to_string(port));

        // 连接到服务器
        try {
            auto ep = net::connect(ws.next_layer(), results);
            std::cout << "Connected to " << ep.address().to_string() << ":" << ep.port() << std::endl;
        } catch (const std::exception& e) {
            std::cerr << "Connection error: " << e.what() << std::endl;
            return 1;
        }

        // 设置WebSocket选项
        ws.set_option(websocket::stream_base::timeout::suggested(beast::role_type::client));
        
        // 设置读写操作的超时时间
        // ws.next_layer().set_option(boost::asio::socket_base::receive_timeout(std::chrono::seconds(10)));
        // ws.next_layer().set_option(boost::asio::socket_base::send_timeout(std::chrono::seconds(10)));

        // 设置服务器主机名（用于WebSocket握手）
        ws.set_option(websocket::stream_base::decorator(
            [](websocket::request_type& req) {
                req.set(http::field::user_agent, "Boost.Beast WebSocket Client");
            }));

        // 执行WebSocket握手
        try {
            ws.handshake(host, "/");
            std::cout << "WebSocket handshake successful" << std::endl;
        } catch (const std::exception& e) {
            std::cerr << "Handshake error: " << e.what() << std::endl;
            return 1;
        }

        // 先尝试发送一个ping消息测试连接
        std::cout << "Sending ping message to test connection..." << std::endl;
        try {
            json ping_message = {
                {"type", "ping"},
                {"timestamp", std::chrono::system_clock::now().time_since_epoch().count()}
            };
            
            if (!send_message(ws, ping_message)) {
                return 1;
            }
            
            // 等待响应（可选）
            std::string ping_response;
            receive_message(ws, ping_response, 2);
        } catch (const std::exception& e) {
            std::cerr << "Ping error (non-fatal): " << e.what() << std::endl;
            // 继续执行，因为服务器可能不处理ping
        }

        // 创建注册请求的JSON
        json register_request = {
            {"type", "register_user"},
            {"username", "user001"},
            {"password", "user001"}
        };

        // 发送注册请求
        if (!send_message(ws, register_request)) {
            return 1;
        }

        // 等待一小段时间，确保服务器有足够时间处理
        std::this_thread::sleep_for(std::chrono::milliseconds(500));

        // 接收响应
        std::string response_str;
        bool got_response = receive_message(ws, response_str);
        
        if (got_response) {
            // 尝试解析响应为JSON
            try {
                json response = json::parse(response_str);
                std::cout << "Response parsed as JSON:" << std::endl;
                std::cout << response.dump(4) << std::endl;
                
                // 检查响应状态
                if (response.contains("success")) {
                    if (response["success"].get<bool>()) {
                        std::cout << "Registration successful!" << std::endl;
                    } else {
                        std::cout << "Registration failed: " << 
                            (response.contains("message") ? response["message"].get<std::string>() : "Unknown error") 
                            << std::endl;
                    }
                }
            } catch (const json::parse_error& e) {
                std::cout << "Failed to parse response as JSON: " << e.what() << std::endl;
            }
        } else {
            std::cout << "No response received from server." << std::endl;
            std::cout << "Server might have crashed or closed the connection." << std::endl;
            std::cout << "Check server logs for double free error or other issues." << std::endl;
        }

        // 检查连接是否仍然打开
        if (ws.is_open()) {
            // 尝试正常关闭WebSocket连接
            try {
                ws.close(websocket::close_code::normal);
                std::cout << "WebSocket connection closed normally" << std::endl;
            } catch (const std::exception& e) {
                std::cerr << "Close error: " << e.what() << std::endl;
            }
        } else {
            std::cout << "Connection was already closed by the server" << std::endl;
        }

    } catch (const std::exception& e) {
        std::cerr << "Unhandled error: " << e.what() << std::endl;
        return 1;
    }

    return 0;
}
